package chan.project.smartlaweveryman.config;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.google.common.collect.Range;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;

import java.text.SimpleDateFormat;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 支持按 timestamp 字段进行精确插入和范围查询的分片算法。
 * 适配 chat_messages_202301, chat_messages_202302 这类表名格式。
 */
@Slf4j
public class MonthBasedTableShardingAlgorithm implements StandardShardingAlgorithm<String> {

    // 时间格式，用于解析 timestamp 字段
    private static final SimpleDateFormat FORMATTER = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
    private static final SimpleDateFormat YYYYMM = new SimpleDateFormat("yyyyMM");

    /**
     * 处理精确分片（插入）
     */
    @Override
    public String doSharding(Collection<String> availableTargetNames, PreciseShardingValue<String> shardingValue) {
        String timestamp = shardingValue.getValue();
        log.info("分表策略生效");
        if (StrUtil.isEmpty(timestamp)) {
            // 默认插入第一个表（可选）
            return availableTargetNames.stream().findFirst()
                    .orElseThrow(() -> new IllegalArgumentException("No available target table found."));
        }

        try {
            Date date = FORMATTER.parse(timestamp);
            String suffix = YYYYMM.format(date);
            String targetTable = shardingValue.getLogicTableName() + "_" + suffix;
            log.info("分表插入：{}", targetTable);
            if (!availableTargetNames.contains(targetTable)) {
                throw new IllegalArgumentException("目标表 " + targetTable + " 不存在");
            }
            return targetTable;
        } catch (Exception e) {
            throw new IllegalArgumentException("时间格式解析错误: " + timestamp, e);
        }
    }

    /**
     * 处理范围分片（查询）
     */
    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, RangeShardingValue<String> shardingValue) {
        String logicTableName = shardingValue.getLogicTableName();
        Range<String> valueRange = shardingValue.getValueRange();

        try {
            Date lower = FORMATTER.parse(valueRange.lowerEndpoint());
            Date upper = FORMATTER.parse(valueRange.upperEndpoint());

            Set<String> targetTables = new HashSet<>();

            // 收集所有可能涉及的月份
            Calendar calendar = Calendar.getInstance();
            calendar.setTime(lower);

            Date current = calendar.getTime();
            while (!current.after(upper)) {
                String suffix = YYYYMM.format(current);
                targetTables.add(logicTableName + "_" + suffix);

                // 移动到下一个月
                calendar.add(Calendar.MONTH, 1);
                current = calendar.getTime();
            }

            // 保留可用表中实际存在的表
            List<String> result = availableTargetNames.stream()
                    .filter(targetTables::contains)
                    .collect(Collectors.toList());

            System.out.println("Sharding tables: " + JSONUtil.toJsonStr(result));
            return result;
        } catch (Exception e) {
            throw new IllegalArgumentException("时间范围解析错误", e);
        }
    }
}
